{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 数据集\n",
    "\n",
    "千帆 Python SDK 支持用户使用 `Dataset` 类在本地进行数据集的创建、删除、查看等操作，然后使用 `DataSource` 类，实现数据集的对不同平台的导入导出功能。\n",
    "\n",
    "接下来，本教程将展示如何在本地导入一个数据集，进行简单的数据集处理与查看，然后分别导出到本地文件以及千帆平台。\n",
    "\n",
    "# 准备工作\n",
    "\n",
    "在开始我们的教程之前，请先升级 SDK 版本至 0.2.8，设置好 Console AK 以及 Console SK 两个环境变量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "pip install -U \"qianfan>=0.2.8\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['QIANFAN_ACCESS_KEY'] = 'YOUR_QIANFAN_ACCESS_KEY'\n",
    "os.environ['QIANFAN_SECRET_KEY'] = 'YOUR_QIANFAN_SECRET_KEY'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# 创建数据集\n",
    "\n",
    "用户可以使用 `Dataset.load()` 方法从一个本地文件创建数据集。该方法只需要传入一个文件路径，即可自动创建一个数据集。\n",
    "\n",
    "下面导入了一个遵循千帆平台有标注对话数据集格式的本地文件。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qianfan.dataset import Dataset\n",
    "\n",
    "ds = Dataset.load(data_file=\"data_file/file.jsonl\", organize_data_as_group=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 数据集处理\n",
    "\n",
    "## 检视数据集\n",
    "\n",
    "如果想要检视这个数据集，我们可以选择调用 `dataset` 对象的 `list()` 方法，该方法会返回包含数据集所有数据的 Python list 对象。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ds.list())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到，该数据集包含 3 组数据，每一组数据都包含一个 `prompt` 字段，表示输入给大模型的 `prompt`，`response` 字段表示大模型的期望输出。上述两个字段的名称、含义以及内容与千帆平台的有标注对话数据集格式对齐，但除此以外，数据集还包含了一个 `_group` 字段。这个字段是由 SDK 在本地中自动生成的，用于区分不同数据集的组。这样的格式更加适合用户在本地以行为单位进行操作。\n",
    "\n",
    "如果用户想要转换成与千帆平台有标注对话数据集格式一致的二维表，可以使用 `dataset` 的 `pack()` 方法，或者在 `load` 时不设置 `organize_data_as_group` 参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.pack()\n",
    "print(ds.list())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此时，整个表中就只剩下由数据所组成的列表，列中的元素对应千帆平台中的每一条数据组。\n",
    "\n",
    "用户还可以选择使用 `unpack()` 方法，将上述格式中的数据重新展开，还原成二维表。为了后续处理方便，我们选择调用 `dataset` 对象的 `unpack()` 方法。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.unpack()\n",
    "print(ds.list())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据集修改\n",
    "\n",
    "`Dataset` 对象中提供了多种可以修改数据集的方法，包括 `append()`，`insert()`， `map()`，`filter()` 等。\n",
    "\n",
    "下面我们分别针对两种格式的数据集，使用这些方法对数据集进行一定程度上的修改。\n",
    "### 字典格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "# 首先创建一个副本引用\n",
    "dict_ds = deepcopy(ds)\n",
    "\n",
    "# 在数据集中间插入一条数据\n",
    "dict_ds.insert({\"prompt\": \"请根据下面的新闻生成摘要, 内容如下:周末采摘杨桃助农，游玩澄迈养生农庄自驾活动，快来报名~1月31日早上8点半出发哦，前往澄迈杨桃园采摘，入园免费，漫山遍野的杨桃和你亲密接触（带走的按2.5元/斤）。报名方式：拨136-3755-3497报名，或戳链接\\n生成摘要如下:\", \"response\": [[\"这个周末自驾游去摘杨桃啦！\"]]}, 1)\n",
    "\n",
    "# 在数据集尾部追加一条数据\n",
    "dict_ds.append({\"prompt\": \"请根据下面的新闻生成摘要, 内容如下:下周最高法第一巡回法庭将开庭审理第一宗案件，最高法巡回法庭历史上的首宗案件，有怎样的特别之处？巡回法庭将首次实行主审法官、合议庭办案负责制，这意味着第一宗案件主审法官将敲响这一制度历史上的“第一槌”\\n生成摘要如下:\", \"response\": [[\"最高法第一巡回法庭将开审首案：跨省买卖合同纠纷\"]]})\n",
    "\n",
    "# 批量追加数据\n",
    "dict_ds.append([\n",
    "    {\"prompt\": \"请根据下面的新闻生成摘要, 内容如下:下周最高法第一巡回法庭将开庭审理第一宗案件，最高法巡回法庭历史上的首宗案件，有怎样的特别之处？巡回法庭将首次实行主审法官、合议庭办案负责制，这意味着第一宗案件主审法官将敲响这一制度历史上的“第一槌”\\n生成摘要如下:\", \"response\": [[\"最高法第一巡回法庭将开审首案：跨省买卖合同纠纷\"]]},\n",
    "    {\"prompt\": \"请根据下面的新闻生成摘要, 内容如下:下周最高法第一巡回法庭将开庭审理第一宗案件，最高法巡回法庭历史上的首宗案件，有怎样的特别之处？巡回法庭将首次实行主审法官、合议庭办案负责制，这意味着第一宗案件主审法官将敲响这一制度历史上的“第一槌”\\n生成摘要如下:\", \"response\": [[\"最高法第一巡回法庭将开审首案：跨省买卖合同纠纷\"]]}\n",
    "])\n",
    "\n",
    "from typing import Any, Dict\n",
    "\n",
    "# 对数据集进行映射处理\n",
    "def map_func(row: Dict[str, Any]) -> Dict[str, Any]:\n",
    "    \"\"\"将 prompt 中的半角冒号替换为全角冒号\"\"\"\n",
    "    prompt = row[\"prompt\"]\n",
    "    row[\"prompt\"] = prompt.replace(\":\", \"：\")\n",
    "    return row\n",
    "\n",
    "# 过滤掉那些 response 过于简短的数据\n",
    "def filter_func(row: Dict[str, Any]) -> bool:\n",
    "    \"\"\"过滤掉 response 中文本长度小于 13 的数据\"\"\"\n",
    "    return len(row[\"response\"][0][0]) > 12\n",
    "\n",
    "dict_ds = dict_ds.map(map_func).filter(filter_func)\n",
    "print(dict_ds.list())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数组格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "from typing import List\n",
    "\n",
    "# 首先创建一个副本引用\n",
    "list_ds = deepcopy(ds)\n",
    "\n",
    "# 并且转换为数组格式\n",
    "list_ds.pack()\n",
    "\n",
    "# 在数据集中间插入一条数据\n",
    "list_ds.insert([{\"prompt\": \"请根据下面的新闻生成摘要, 内容如下:周末采摘杨桃助农，游玩澄迈养生农庄自驾活动，快来报名~1月31日早上8点半出发哦，前往澄迈杨桃园采摘，入园免费，漫山遍野的杨桃和你亲密接触（带走的按2.5元/斤）。报名方式：拨136-3755-3497报名，或戳链接\\n生成摘要如下:\", \"response\": [[\"这个周末自驾游去摘杨桃啦！\"]]}], 1)\n",
    "\n",
    "# 在数据集尾部追加一条数据\n",
    "# 直接插入一个字典会被特殊处理成插入一个数组，其中包括该字典对象\n",
    "list_ds.append({\"prompt\": \"请根据下面的新闻生成摘要, 内容如下:下周最高法第一巡回法庭将开庭审理第一宗案件，最高法巡回法庭历史上的首宗案件，有怎样的特别之处？巡回法庭将首次实行主审法官、合议庭办案负责制，这意味着第一宗案件主审法官将敲响这一制度历史上的“第一槌”\\n生成摘要如下:\", \"response\": [[\"最高法第一巡回法庭将开审首案：跨省买卖合同纠纷\"]]})\n",
    "\n",
    "# 批量追加数据\n",
    "list_ds.append(\n",
    "    [\n",
    "        [{\"prompt\": \"请根据下面的新闻生成摘要, 内容如下:下周最高法第一巡回法庭将开庭审理第一宗案件，最高法巡回法庭历史上的首宗案件，有怎样的特别之处？巡回法庭将首次实行主审法官、合议庭办案负责制，这意味着第一宗案件主审法官将敲响这一制度历史上的“第一槌”\\n生成摘要如下:\", \"response\": [[\"最高法第一巡回法庭将开审首案：跨省买卖合同纠纷\"]]}],\n",
    "        [{\"prompt\": \"请根据下面的新闻生成摘要, 内容如下:下周最高法第一巡回法庭将开庭审理第一宗案件，最高法巡回法庭历史上的首宗案件，有怎样的特别之处？巡回法庭将首次实行主审法官、合议庭办案负责制，这意味着第一宗案件主审法官将敲响这一制度历史上的“第一槌”\\n生成摘要如下:\", \"response\": [[\"最高法第一巡回法庭将开审首案：跨省买卖合同纠纷\"]]}]\n",
    "    ]\n",
    ")\n",
    "\n",
    "from typing import Any, Dict\n",
    "\n",
    "# 对列表行进行映射处理\n",
    "def map_func(row: List[Dict[str, Any]]) -> List[Dict[str, Any]]:\n",
    "    \"\"\"将 prompt 中的半角冒号替换为全角冒号\"\"\"\n",
    "    for obj in row:\n",
    "        prompt = obj[\"prompt\"]\n",
    "        obj[\"prompt\"] = prompt.replace(\":\", \"：\")\n",
    "    \n",
    "    return row\n",
    "\n",
    "dedup_list = []\n",
    "\n",
    "# 过滤掉那些重复的数据\n",
    "def filter_func(row: List[Dict[str, Any]]) -> bool:\n",
    "    \"\"\"过滤重复的数据\"\"\"\n",
    "    if row in dedup_list:\n",
    "        return False\n",
    "    dedup_list.append(row)\n",
    "    return True\n",
    "\n",
    "list_ds = list_ds.map(map_func).filter(filter_func)\n",
    "print(list_ds.list())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 导出数据集\n",
    "\n",
    "在完成上述操作时候，我们就可以将数据集导出到本地文件或者千帆平台。\n",
    "\n",
    "但首先，我们先将改动同步到 `ds` 对象上"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = list_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## 导出到本地文件\n",
    "\n",
    "用户只需要调用 `dataset` 对象的 `save()` 方法，传入一个文件路径到 `data_file` 参数，即可将数据集导出到本地文件。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.save(data_file=\"data_file/post_processed_file.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`save()` 方法还支持传入 `schema` 参数，用于对导出的数据集进行格式校验，以便用户后续自行处理，如手动将数据集文件导入到千帆平台中。 \n",
    "\n",
    "SDK 内集成的 schema 保存于 qianfan.dataset.schema 中，里面包含了与千帆平台相关的所有 schema 实现类。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qianfan.dataset.schema import QianfanNonSortedConversation\n",
    "\n",
    "ds.save(data_file=\"data_file/post_processed_file.jsonl\", schema=QianfanNonSortedConversation())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 导出到千帆平台\n",
    "\n",
    "导出到千帆平台分为两种情况：\n",
    "1. 导出到已经存在的平台数据集中。这种情况下，需要用户自行确认平台数据集的格式是否与本地数据集所要求的格式一致。否则导出时的内置校验会失败。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "your_dataset_id = \"42\"\n",
    "ds.save(qianfan_dataset_id=your_dataset_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果上传的数据集并非使用私有 BOS 作为存储，则需要设置 BOS 配置用户上传"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "your_dataset_id = \"42\"\n",
    "bos_bucket_name = \"your_bos_bucket\"\n",
    "bos_bucket_file_path = \"/your_bos_bucket_file_path/\"\n",
    "bos_bucket_region = \"your_bos_region\"\n",
    "\n",
    "ds.save(\n",
    "    qianfan_dataset_id=your_dataset_id,\n",
    "    sup_storage_id=bos_bucket_name, \n",
    "    sup_storage_path=bos_bucket_file_path, \n",
    "    sup_storage_region=bos_bucket_region,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. 在导出的同时创建一个全新的平台数据集。这种情况下，又有两种方式：\n",
    "\n",
    "    A. 使用 `QianfanDataSource` 提供的 `create_bare_dataset()` 方法，在本地创建一个数据源的同时，在千帆平台上创建一个对应的数据源\n",
    "\n",
    "    B. 在 `dataset` 的 `save()` 方法中传入 `qianfan_dataset_create_args` 参数，该参数是一个字典对象，用于指定在千帆平台上创建数据集时所需要的信息，以便在导出时自动创建一个新的数据集。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 对于方式 A:\n",
    "from qianfan.dataset import (\n",
    "    DataTemplateType,\n",
    "    DataStorageType,\n",
    ")\n",
    "\n",
    "from qianfan.dataset.data_source import QianfanDataSource\n",
    "\n",
    "qianfan_dataset_name = \"your_dataset_name\"\n",
    "\n",
    "bos_bucket_name = \"your_bos_bucket\"\n",
    "bos_bucket_file_path = \"/your_bos_bucket_file_path/\"\n",
    "\n",
    "qianfan_data_source = QianfanDataSource.create_bare_dataset(\n",
    "    name=qianfan_dataset_name,\n",
    "    template_type=DataTemplateType.NonSortedConversation,\n",
    "    storage_type=DataStorageType.PrivateBos,\n",
    "    storage_id=bos_bucket_name,\n",
    "    storage_path=bos_bucket_file_path,\n",
    ")\n",
    "\n",
    "ds.save(qianfan_data_source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 对于方式 B:\n",
    "\n",
    "from qianfan.dataset import (\n",
    "    DataTemplateType,\n",
    "    DataStorageType,\n",
    ")\n",
    "\n",
    "qianfan_dataset_name = \"test_name_1213\"\n",
    "\n",
    "bos_bucket_name = \"test-bucket-of-qianfan\"\n",
    "bos_bucket_file_path = \"/test_path/\"\n",
    "\n",
    "ds.save(\n",
    "    qianfan_dataset_create_args={\n",
    "        \"name\": qianfan_dataset_name,\n",
    "        \"template_type\": DataTemplateType.NonSortedConversation,\n",
    "        \"storage_type\": DataStorageType.PrivateBos,\n",
    "        \"storage_id\": bos_bucket_name,\n",
    "        \"storage_path\": bos_bucket_file_path,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在导出完成后，用户即可以在千帆平台上查看到导出的数据集内容。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "58f7cb64c3a06383b7f18d2a11305edccbad427293a2b4afa7abe8bfc810d4bb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
